import cv2
import mediapipe as mp
import numpy as np

# 初始化 MediaPipe
mp_pose = mp.solutions.pose
pose = mp_pose.Pose()
mp_drawing = mp.solutions.drawing_utils

# 读取输入图片
image_path = "./zishi.png"  # 替换为实际路径
image = cv2.imread(image_path)
cv2.imshow("yuanshi", image)
image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

# 检测姿势
results = pose.process(image_rgb)

# 创建空白背景图
output_image = np.zeros(image.shape, dtype=np.uint8)

if results.pose_landmarks:
    # 提取关键点坐标
    landmarks = results.pose_landmarks.landmark
    height, width, _ = image.shape

    # 绘制骨架图
    for landmark in landmarks:
        x = int(landmark.x * width)
        y = int(landmark.y * height)
        # 绘制关节为圆形
        cv2.circle(output_image, (x, y), radius=10, color=(255, 255, 255), thickness=-1)

    # 手动绘制关键点连接
    connections = mp_pose.POSE_CONNECTIONS
    for connection in connections:
        start_idx, end_idx = connection
        start = landmarks[start_idx]
        end = landmarks[end_idx]
        x1, y1 = int(start.x * width), int(start.y * height)
        x2, y2 = int(end.x * width), int(end.y * height)
        cv2.line(output_image, (x1, y1), (x2, y2), color=(255, 255, 255), thickness=5)

# 保存输出图片
output_path = "./output_skeleton.png"
cv2.imwrite(output_path, output_image)

# 显示结果
cv2.imshow("Skeleton", output_image)
cv2.waitKey(0)
cv2.destroyAllWindows()
